import torch
import torch.utils.data as data
import os, math, random
from os import path as osp
import numpy as np
from glob import glob
from pathlib import Path
from mynn.utils.registry import DATASET_REGISTRY
from mynn.utils import FileClient, imfrombytes, flowread, img2tensor
from PIL import Image


class StaticRandomCrop(object):
    def __init__(self, image_size, crop_size, is_random=True):
        self.th, self.tw = crop_size
        h, w = image_size
        self.h1 = random.randint(0, h - self.th)
        self.w1 = random.randint(0, w - self.tw)
        if is_random is not True:
            self.h1 = 0
            self.w1 = 0

    def __call__(self, img):
        return img[self.h1:(self.h1 + self.th), self.w1:(self.w1 + self.tw), :]


class StaticCenterCrop(object):
    def __init__(self, image_size, crop_size):
        self.th, self.tw = crop_size
        self.h, self.w = image_size

    def __call__(self, img):
        return img[(self.h - self.th) // 2:(self.h + self.th) // 2, (self.w - self.tw) // 2:(self.w + self.tw) // 2, :]


@DATASET_REGISTRY.register()
class Sintel(data.Dataset):
    def __init__(self, dataset_opt, phase='train'):
        _phase = ['train', 'val', 'test']
        assert phase in _phase, f"Wrong phase {phase}, supported choices are {_phase}"

        self.dataset_opt = dataset_opt
        self.phase = phase
        self.io_backend_opt = dataset_opt['io_backend']
        self.flow_root = Path(dataset_opt['flow_root'])
        self.image_root = Path(dataset_opt['image_root'])
        self.keys = []
        self.file_client = None
        self.clip_length = {}
        with open(dataset_opt['meta_info_file'], 'r') as fin:
            for line in fin:
                folder, frame_num, _ = line.split(' ')
                self.clip_length[f'{folder}'] = int(frame_num)
                self.keys.extend([f'{folder}/frame_{i+1:04d}' for i in range(int(frame_num))])

        if phase == 'train' or phase == 'val':
            if dataset_opt['val_partition'] == 'Sintel2':
                val_partition = ['ambush_7', 'market_6']
            else:
                raise ValueError(f'Wrong validation partition {dataset_opt["val_partition"]}.'
                                 f"Supported choices are ['Sintel2'].")

        if phase == 'train':
            self.keys = [v for v in self.keys if v.split('/')[0] not in val_partition]
            self.crop_size = dataset_opt['crop_size']
        elif phase == 'val':
            self.keys = [v for v in self.keys if v.split('/')[0] in val_partition]
            self.crop_size = None
        elif phase == 'test':
            self.crop_size = None

    def __getitem__(self, index):
        if self.file_client is None:
            self.file_client = FileClient()

        index = index % (len(self.keys))

        key = self.keys[index]

        clip_name, frame_name = key.split('/')
        clip_length = self.clip_length[clip_name]
        ref_frame_idx = int(frame_name[-4:])
        supp_frame_idx = ref_frame_idx + 1
        while (supp_frame_idx > clip_length) and self.phase != 'test':
            ref_frame_idx = random.randint(1, clip_length)
            supp_frame_idx = ref_frame_idx + 1

        ref_frame_name = f'frame_{ref_frame_idx:04d}'
        supp_frame_name = f'frame_{supp_frame_idx:04d}'

        ref_frame_path = self.image_root / clip_name / f'{ref_frame_name}.png'
        supp_frame_path = self.image_root / clip_name / f'{supp_frame_name}.png'
        flow_path = self.flow_root / clip_name / f'{ref_frame_name}.flo'

        ref_bytes = self.file_client.get(ref_frame_path)
        ref = imfrombytes(ref_bytes, float32=True)
        supp_bytes = self.file_client.get(supp_frame_path)
        supp = imfrombytes(supp_bytes, float32=True)

        flow = flowread(flow_path=flow_path)

        image_size = ref.shape[:2]

        if self.crop_size is not None:
            cropper = StaticRandomCrop(image_size, self.crop_size)
        else:
            cropper = StaticRandomCrop(image_size, (image_size[0] // 16 * 16, image_size[1] // 16 * 16),
                                       is_random=False)

        ref = img2tensor(cropper(ref))
        supp = img2tensor(cropper(supp))
        flow = img2tensor(cropper(flow))

        return {'ref': ref, 'supp': supp, 'flow': flow, 'key': key}

    def __len__(self):
        return len(self.keys)
